import pandas as pd
import matplotlib.pyplot as plt

# ==================== 配置参数 ====================
INPUT_FILE = "ori_message.xlsx"  # 原始数据文件路径
OUTPUT_FILE = "monthly_average_data.xlsx"  # 输出文件路径


# ==================== 数据读取与预处理 ====================
def process_data():
    excel_file = pd.ExcelFile(INPUT_FILE)
    years = ['2016', '2017', '2018', '2019', '2020', '2021']
    all_monthly_data = []
    for year in years:
        df = excel_file.parse(year)
        df['年'] = df['年'].ffill().astype(int)
        df['月'] = df['月'].ffill().astype(int)
        grouped = df.groupby(['年', '月']).agg({
            '流量(m3/s)': 'mean',
            '含沙量(kg/m3) ': 'mean'
        }).reset_index()
        grouped['排沙量(kg/s)'] = grouped['流量(m3/s)'] * grouped['含沙量(kg/m3) ']
        grouped = grouped.round({
            '流量(m3/s)': 2,
            '含沙量(kg/m3) ': 4,
            '排沙量(kg/s)': 2
        })
        all_monthly_data.append(grouped)
    result = pd.concat(all_monthly_data, ignore_index=True)
    result = result.sort_values(by=['年', '月']).reset_index(drop=True)
    result = result.rename(columns={'含沙量(kg/m3) ': '含沙量(kg/m3)'})
    result.to_excel(OUTPUT_FILE, index=False)
    return result


# ==================== 绘制5-10月排沙量走势图 ====================
def plot_5_10_sediment(data):
    # 筛选5-10月数据
    filtered_data = data[data['月'].between(5, 10)]

    # 定义年份顺序和颜色（与图30示例对应）
    years_order = [2016, 2017, 2018, 2019, 2020, 2021]
    colors = ['#0000FF', '#FF0000', '#808080', '#FFFF00', '#1E90FF', '#008000']  # 对应图中颜色

    plt.figure(figsize=(10, 6))

    for i, year in enumerate(years_order):
        year_data = filtered_data[filtered_data['年'] == year]
        months = year_data['月'].map({5: '5月', 6: '6月', 7: '7月', 8: '8月', 9: '9月', 10: '10月'})
        plt.plot(months, year_data['排沙量(kg/s)'], marker='o', linestyle='-',
                 label=str(year), color=colors[i], linewidth=1.5)

    plt.title('六年中5-10月的月均排沙量走势图', fontsize=12)
    plt.xlabel('')
    plt.ylabel('月平均排沙量 (吨)', fontsize=10)  # 假设单位为吨，根据实际情况调整
    plt.xticks(rotation=0)
    plt.legend(loc='best', title='年份', bbox_to_anchor=(1, 1))
    plt.grid(linestyle='--', alpha=0.5, axis='y')
    plt.tight_layout(rect=[0, 0, 0.85, 1])

    plt.savefig('图30.六年中5-10月的月均排沙量走势图.png', dpi=300, bbox_inches='tight')
    plt.show()


# ==================== 主程序 ====================
if __name__ == "__main__":
    processed_data = process_data()
    plot_5_10_sediment(processed_data)